{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "执行测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from test import test_all"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test_all()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def warp_boxes(boxes, theta, img_size):\n",
    "    \"\"\"\n",
    "    boxes: Tensor(N,4)  YOLO [cx,cy,w,h] 归一化\n",
    "    theta: Tensor(B,2,3) 仿射矩阵\n",
    "    return: 同样形状的框，但已 warp\n",
    "    \"\"\"\n",
    "    # 先转绝对坐标\n",
    "    H, W = img_size\n",
    "    cx, cy, w, h = boxes.T\n",
    "    x1, y1 = (cx - w / 2) * W, (cy - h / 2) * H\n",
    "    x2, y2 = (cx + w / 2) * W, (cy + h / 2) * H\n",
    "    pts = torch.stack(\n",
    "        [\n",
    "            torch.stack([x1, y1, torch.ones_like(x1)], dim=1),\n",
    "            torch.stack([x2, y2, torch.ones_like(x2)], dim=1),\n",
    "        ],\n",
    "        dim=1,\n",
    "    )  # (N,2,3)\n",
    "\n",
    "    # 仿射变换\n",
    "    pts = pts @ theta.transpose(1, 2)  # (N,2,2)\n",
    "    pts = pts / torch.tensor([W, H], device=pts.device)  # 重新归一化\n",
    "\n",
    "    # 再变回 YOLO\n",
    "    new_cx = (pts[:, 0, 0] + pts[:, 1, 0]) / 2\n",
    "    new_cy = (pts[:, 0, 1] + pts[:, 1, 1]) / 2\n",
    "    new_w = pts[:, 1, 0] - pts[:, 0, 0]\n",
    "    new_h = pts[:, 1, 1] - pts[:, 0, 1]\n",
    "    return torch.stack([new_cx, new_cy, new_w, new_h], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def yolo2xyxy(boxes):\n",
    "    \"\"\"\n",
    "    boxes: Tensor(N,4)  归一化 YOLO [cx, cy, w, h]\n",
    "    return: Tensor(N,4) [x1, y1, x2, y2]  仍归一化\n",
    "    \"\"\"\n",
    "    cx, cy, w, h = boxes.T\n",
    "    x1 = cx - w / 2\n",
    "    y1 = cy - h / 2\n",
    "    x2 = cx + w / 2\n",
    "    y2 = cy + h / 2\n",
    "    return torch.stack([x1, y1, x2, y2], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "# Fast R-CNN 检测模型\n",
    "from torchvision.models.detection import fasterrcnn_resnet50_fpn\n",
    "from torchvision.models import resnet18\n",
    "from STN import STN, SpatialTransformer\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision.models.detection import fasterrcnn_resnet50_fpn\n",
    "from torchvision.models.resnet import resnet50\n",
    "from torchvision.models.detection.backbone_utils import (\n",
    "    _resnet_fpn_extractor,\n",
    "    _validate_trainable_layers,\n",
    ")\n",
    "from torchvision.ops import misc as misc_nn_ops\n",
    "from DamageDetectionModel import STN_FasterRCNN\n",
    "from DataLoader import FLIR_ImagePair_Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_CLASSES = 3  # 3 类目标检测\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "root = \"/root/MyCode/infrared-image-damage-detection-model/dataset/FLIR-align-3class\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(batch):\n",
    "    # batch: list[dict]\n",
    "    vis = [b[\"vis\"] for b in batch]  # (3,H,W)\n",
    "    ir = [b[\"ir\"] for b in batch]  # (1,H,W)\n",
    "    label = [b[\"label\"] for b in batch]  # (N,5)\n",
    "    return vis, ir, label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds = FLIR_ImagePair_Dataset(root, split=\"train\")\n",
    "val_ds = FLIR_ImagePair_Dataset(root, split=\"test\")\n",
    "train_loader = DataLoader(\n",
    "    train_ds,\n",
    "    batch_size=4,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    collate_fn=collate,\n",
    ")\n",
    "\n",
    "val_loader = DataLoader(\n",
    "    val_ds,\n",
    "    batch_size=4,\n",
    "    shuffle=False,\n",
    "    num_workers=4,\n",
    "    collate_fn=collate,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = STN_FasterRCNN(num_classes=3 + 1).to(device)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "CUDA out of memory. Tried to allocate 40.00 MiB (GPU 0; 11.91 GiB total capacity; 4.93 GiB already allocated; 25.62 MiB free; 5.05 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "\u001b[1;32m/root/MyCode/infrared-image-damage-detection-model/main.ipynb Cell 12\u001b[0m line \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/main.ipynb#X10sdnNjb2RlLXJlbW90ZQ%3D%3D?line=5'>6</a>\u001b[0m vis \u001b[39m=\u001b[39m [v\u001b[39m.\u001b[39mto(device) \u001b[39mfor\u001b[39;00m v \u001b[39min\u001b[39;00m vis]\n\u001b[1;32m      <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/main.ipynb#X10sdnNjb2RlLXJlbW90ZQ%3D%3D?line=6'>7</a>\u001b[0m ir \u001b[39m=\u001b[39m [i\u001b[39m.\u001b[39mto(device) \u001b[39mfor\u001b[39;00m i \u001b[39min\u001b[39;00m ir]\n\u001b[0;32m----> <a href='vscode-notebook-cell://localhost:8080/root/MyCode/infrared-image-damage-detection-model/main.ipynb#X10sdnNjb2RlLXJlbW90ZQ%3D%3D?line=7'>8</a>\u001b[0m losses \u001b[39m=\u001b[39m model(vis, ir, targets)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/MyCode/infrared-image-damage-detection-model/DamageDetectionModel.py:155\u001b[0m, in \u001b[0;36mSTN_FasterRCNN.forward\u001b[0;34m(self, vis, ir, targets)\u001b[0m\n\u001b[1;32m    151\u001b[0m fused \u001b[39m=\u001b[39m [\u001b[39mself\u001b[39m\u001b[39m.\u001b[39mgtf(v\u001b[39m.\u001b[39munsqueeze(\u001b[39m0\u001b[39m), ia\u001b[39m.\u001b[39munsqueeze(\u001b[39m0\u001b[39m))\u001b[39m.\u001b[39msqueeze(\u001b[39m0\u001b[39m)\n\u001b[1;32m    152\u001b[0m         \u001b[39mfor\u001b[39;00m v, ia \u001b[39min\u001b[39;00m \u001b[39mzip\u001b[39m(vis, ir_aligned)]\n\u001b[1;32m    154\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mtraining:\n\u001b[0;32m--> 155\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mdetector(fused, targets)\n\u001b[1;32m    156\u001b[0m \u001b[39melse\u001b[39;00m:\n\u001b[1;32m    157\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdetector(fused)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torchvision/models/detection/generalized_rcnn.py:101\u001b[0m, in \u001b[0;36mGeneralizedRCNN.forward\u001b[0;34m(self, images, targets)\u001b[0m\n\u001b[1;32m     94\u001b[0m             degen_bb: List[\u001b[39mfloat\u001b[39m] \u001b[39m=\u001b[39m boxes[bb_idx]\u001b[39m.\u001b[39mtolist()\n\u001b[1;32m     95\u001b[0m             torch\u001b[39m.\u001b[39m_assert(\n\u001b[1;32m     96\u001b[0m                 \u001b[39mFalse\u001b[39;00m,\n\u001b[1;32m     97\u001b[0m                 \u001b[39m\"\u001b[39m\u001b[39mAll bounding boxes should have positive height and width.\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m     98\u001b[0m                 \u001b[39mf\u001b[39m\u001b[39m\"\u001b[39m\u001b[39m Found invalid box \u001b[39m\u001b[39m{\u001b[39;00mdegen_bb\u001b[39m}\u001b[39;00m\u001b[39m for target at index \u001b[39m\u001b[39m{\u001b[39;00mtarget_idx\u001b[39m}\u001b[39;00m\u001b[39m.\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m     99\u001b[0m             )\n\u001b[0;32m--> 101\u001b[0m features \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbackbone(images\u001b[39m.\u001b[39;49mtensors)\n\u001b[1;32m    102\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39misinstance\u001b[39m(features, torch\u001b[39m.\u001b[39mTensor):\n\u001b[1;32m    103\u001b[0m     features \u001b[39m=\u001b[39m OrderedDict([(\u001b[39m\"\u001b[39m\u001b[39m0\u001b[39m\u001b[39m\"\u001b[39m, features)])\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torchvision/models/detection/backbone_utils.py:57\u001b[0m, in \u001b[0;36mBackboneWithFPN.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     56\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, x: Tensor) \u001b[39m-\u001b[39m\u001b[39m>\u001b[39m Dict[\u001b[39mstr\u001b[39m, Tensor]:\n\u001b[0;32m---> 57\u001b[0m     x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbody(x)\n\u001b[1;32m     58\u001b[0m     x \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mfpn(x)\n\u001b[1;32m     59\u001b[0m     \u001b[39mreturn\u001b[39;00m x\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torchvision/models/_utils.py:69\u001b[0m, in \u001b[0;36mIntermediateLayerGetter.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     67\u001b[0m out \u001b[39m=\u001b[39m OrderedDict()\n\u001b[1;32m     68\u001b[0m \u001b[39mfor\u001b[39;00m name, module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mitems():\n\u001b[0;32m---> 69\u001b[0m     x \u001b[39m=\u001b[39m module(x)\n\u001b[1;32m     70\u001b[0m     \u001b[39mif\u001b[39;00m name \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mreturn_layers:\n\u001b[1;32m     71\u001b[0m         out_name \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mreturn_layers[name]\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/container.py:139\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    137\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m):\n\u001b[1;32m    138\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m:\n\u001b[0;32m--> 139\u001b[0m         \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m module(\u001b[39minput\u001b[39;49m)\n\u001b[1;32m    140\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39minput\u001b[39m\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torchvision/models/resnet.py:155\u001b[0m, in \u001b[0;36mBottleneck.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    152\u001b[0m out \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mrelu(out)\n\u001b[1;32m    154\u001b[0m out \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mconv3(out)\n\u001b[0;32m--> 155\u001b[0m out \u001b[39m=\u001b[39m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbn3(out)\n\u001b[1;32m    157\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdownsample \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    158\u001b[0m     identity \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mdownsample(x)\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/opt/anaconda3/envs/develop/lib/python3.9/site-packages/torchvision/ops/misc.py:62\u001b[0m, in \u001b[0;36mFrozenBatchNorm2d.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     60\u001b[0m scale \u001b[39m=\u001b[39m w \u001b[39m*\u001b[39m (rv \u001b[39m+\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39meps)\u001b[39m.\u001b[39mrsqrt()\n\u001b[1;32m     61\u001b[0m bias \u001b[39m=\u001b[39m b \u001b[39m-\u001b[39m rm \u001b[39m*\u001b[39m scale\n\u001b[0;32m---> 62\u001b[0m \u001b[39mreturn\u001b[39;00m x \u001b[39m*\u001b[39;49m scale \u001b[39m+\u001b[39;49m bias\n",
      "\u001b[0;31mRuntimeError\u001b[0m: CUDA out of memory. Tried to allocate 40.00 MiB (GPU 0; 11.91 GiB total capacity; 4.93 GiB already allocated; 25.62 MiB free; 5.05 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF"
     ]
    }
   ],
   "source": [
    "for vis, ir, labels in train_loader:\n",
    "    targets = [\n",
    "        {\"boxes\": yolo2xyxy(l[:, 1:]).to(device), \"labels\": l[:, 0].long().to(device)}\n",
    "        for l in labels\n",
    "    ]\n",
    "    vis = [v.to(device) for v in vis]\n",
    "    ir = [i.to(device) for i in ir]\n",
    "    losses = model(vis, ir, targets)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
